import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import codecs


def PlotPef_on_FewSet(log_dir):
    log_files = [fname for fname in os.listdir(log_dir) if ".log" in fname]
    for filename in log_files:
        log_file = os.path.join(log_dir, filename)
        fewshot_losses = []
        fewshot_accs = []
        tmp_losses, tmp_accs = [], []
        with codecs.open(log_file, 'r', encoding= u'utf-8',errors='ignore') as fr:
            for line in fr:
                if "-------> few shot data list ------>" in line:
                    if len(tmp_accs) != 0 and len(tmp_losses) != 0:
                        fewshot_losses.append(np.mean(tmp_losses))
                        fewshot_accs.append(np.mean(tmp_accs))
                    tmp_losses, tmp_accs = [], []
                if "####Few Shot" in line:
                    loss, acc = line.split("=")[1].split("/")
                    tmp_losses.append(float(loss))
                    tmp_accs.append(float(acc))
        if len(fewshot_accs) == 0 and len(fewshot_losses)==0:
            fewshot_accs, fewshot_losses = tmp_accs, tmp_losses
        steps = list(range(len(fewshot_accs)))
        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_accs, label="accuracy")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Acc.png")

        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_losses, label="loss")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Loss.png")

def Plot_TrainingSet(log_file):

    slc_accs, train_accs, pseu_accs = [], [], []
    with codecs.open(log_file, 'r', encoding='utf-8', errors="ignore") as fr:
        for line in fr:
            if "###Accuracy On selected instancees" in line:
                acc = float(
                    line.strip("###Accuracy On selected instancees (").split(",")[0]
                )
                slc_accs.append(acc)
            if "##Accuracy On pseaudo instances" in line:
                acc = float(
                    line.strip("###Accuracy On pseaudo instancees (").split(",")[0]
                )
                pseu_accs.append(acc)

            if "###Accuracy On training instances" in line:
                acc = float(
                    line.strip("###Accuracy On training instancees (").split(",")[0]
                )
                train_accs.append(acc)

    slc_accs1 = [item for i, item in enumerate(slc_accs) if i % 2 == 0]
    slc_accs2 = [item for i, item in enumerate(slc_accs) if i % 2 == 1]
    train_accs1 = [item for i, item in enumerate(train_accs) if i % 2 == 0]
    train_accs2 = [item for i, item in enumerate(train_accs) if i % 2 == 1]
    pseu_accs1 = [item for i, item in enumerate(pseu_accs) if i % 2 == 0]
    pseu_accs2 = [item for i, item in enumerate(pseu_accs) if i % 2 == 1]

    plt.cla()
    plt.title("Accuracy On Selected Instances")
    plt.xlabel("steps")
    plt.plot(range(len(slc_accs1)), slc_accs1, label="model 1")
    plt.plot(range(len(slc_accs2)), slc_accs2, label="model 2")
    plt.legend()
    plt.savefig(f"./Selected_Acc.png")

    plt.cla()
    plt.title("Accuracy On Pseudo Instances")
    plt.xlabel("steps")
    plt.plot(range(len(pseu_accs1)), pseu_accs1, label="model 1")
    plt.plot(range(len(pseu_accs2)), pseu_accs2, label="model 2")
    plt.legend()
    plt.savefig(f"./Pseudo_Acc.png")

    plt.cla()
    plt.title("Accuracy On Train Instances")
    plt.xlabel("steps")
    plt.plot(range(len(train_accs1)), train_accs1, label="model 1")
    plt.plot(range(len(train_accs2)), train_accs2, label="model 2")
    plt.legend()
    plt.savefig(f"./train_Acc.png")

Plot_TrainingSet(sys.argv[1])